import os
import re
import json
from typing import List
from fastapi import FastAPI
from pydantic import BaseModel
from openai import OpenAI
import uvicorn

# -----------------------
# API 配置（全局）
# -----------------------
API_KEY = os.environ.get("ARK_API_KEY", "sk-U4BpSvt3NnkGhZpYtZtjmOjWR0YKSkScXVLBk6QIYHaGgaQC")
BASE_URL = os.environ.get("ARK_API_BASE_URL", "http://ai.docker.tcl.com/huizhou/inference/usr-gqb23upd/inf-7qgwnmfo/v1")  # 例如 https://api.openai.com/v1

client = OpenAI(
    api_key=API_KEY,
    base_url=BASE_URL
)

# -----------------------
# 角色设定
# -----------------------
agent_role_design = '''你是一个熟悉公司经营管理流程的总裁助理，你和你服务的高管已经共事超过5年；
你非常了解经营者们的工作习惯，当他们向你提出了一些关于公司经营状况的问题后，你不仅能快速用关键数据信息回复并且能够猜到他们后续还会有哪些疑问；
现在你要运用你的能力来完成一些特定的工作。
'''

remake_query_prompt_template = '''一位{}的用户还没有想好下面要向你提出什么问题，你需要给他推荐一些他可能感兴趣的提问。
{}

你推荐的问题不能和用户已经提出的问题出现重复，这种重复指的是语义上的重复而不是仅指内容上的重复；你推荐的问题要确保和之前用户提出的问题在同一个经营领域，不能太过跳脱，导致推荐了用户根本不感兴趣的问题。
请按照下面的json格式输出你推荐的提问，记住不要包含任何json以外的内容：
[
  {{
    "qid": "1",
    "content": ""
  }},
  {{
    "qid": "2",
    "content": ""
  }},
  {{
    "qid": "3",
    "content": ""
  }}
]

推荐至少三个提问，但不要超过八个。
'''

# -----------------------
# Agent Memory
# -----------------------
class Agent_Memory:
    def __init__(self):
        self.user_identity: str = ''
        self.query_record: List[str] = []

    def merge_memory(self):
        if not self.query_record:
            return "用户还没有提问过问题。"
        history_content = "用户已经进行了以下提问：\n"
        for i, query in enumerate(self.query_record):
            history_content += "第{}次提问：{}\n".format(i+1, query)
        return history_content

# -----------------------
# 大模型调用函数
# -----------------------
def suggest_agent(memory: Agent_Memory):
    user_id = memory.user_identity
    history = memory.merge_memory()
    response = client.chat.completions.create(
        model="",  
        messages=[
            {"role": "system", "content": agent_role_design},
            {"role": "user", "content": remake_query_prompt_template.format(user_id, history)}
        ],
        extra_body={"enable_think": False},
        stream=False
    )

    content = response.choices[0].message.content

    # 清理思维链
    content = re.sub(r'\<think\>[\s\S]*?\<\/think\>', '', content).strip()
    content = re.sub(r'^(\n+)?```json\n', '', content, count=1)
    content = re.sub(r'\n```$', '', content, count=1)

    try:
        suggest_json = json.loads(content)
    except Exception:
        suggest_json = [{"qid": "0", "content": "解析失败，原始输出：" + content}]

    return suggest_json

# -----------------------
# FastAPI 封装
# -----------------------
app = FastAPI()

class QueryRequest(BaseModel):
    user_identity: str = ""
    user_question: str

@app.post("/suggest")
def suggest_api(req: QueryRequest):
    memory = Agent_Memory()
    memory.user_identity = req.user_identity
    memory.query_record.append(req.user_question)

    suggestions = suggest_agent(memory)
    return {"recommendations": suggestions}

# -----------------------
# main 函数（可单独运行测试）
# -----------------------
def main():
    record = Agent_Memory()
    record.user_identity = "财务经理"

    user_question = "公司今年的现金流情况怎么样？"
    record.query_record.append(user_question)

    suggestions = suggest_agent(record)
    print("用户提问：", user_question)
    print("推荐的问题：", json.dumps(suggestions, ensure_ascii=False, indent=2))

if __name__ == '__main__':
    #main()
    uvicorn.run(app, host="0.0.0.0", port=8000)